from __future__ import division
import sys
import pygame
import auxiliar
from matrix import *
from cell import *
import agente
import random
import math
import util 


BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
DOWN = False
CLEAR = False
cells = []

########################################33
# Geometria de los agentes
RADIO_AGENTES = 40
COLOR_AGENTES = 'blue'
SIZE = 6

########################################33
# velocidad minima e inicial
MIN_VEL = 10.0
INIT_VEL = MIN_VEL
# Radio de vision
VISION = 3*RADIO_AGENTES

# Posicion inicial de los agentes y masa
P0 = (SIZE*60/2,100)
P1 = (P0[0]+RADIO_AGENTES*1,P0[1]-RADIO_AGENTES*1)
P2 = (P0[0]-RADIO_AGENTES*1,P0[1]-RADIO_AGENTES*1)
P3 = (P0[0]+RADIO_AGENTES*2,P0[1]-RADIO_AGENTES*2)
P4 = (P0[0]-RADIO_AGENTES*2,P0[1]-RADIO_AGENTES*2) 
P5 = (P0[0]+RADIO_AGENTES*3,P0[1]-RADIO_AGENTES*1)
P6 = (SIZE/2, SIZE/2)
P7 = (SIZE/4, SIZE/4)
P8 = (SIZE/7, SIZE/7)

m0, m1, m2, m3, m4, m5, m6, m7, m8 = 5, 5, 5, 5, 5, 5, 5, 5, 5

#POS = [[[P0, m0], [P1,m1], [P2,m2], [P3,m3], [P4,m4]]]
INIT_POS = [P0, P1, P2, P3, P4, P5, P6, P7, P8]
INIT_VELO = [[INIT_VEL, 90], [INIT_VEL, 90], [INIT_VEL, 90], [INIT_VEL, 90], [INIT_VEL, 90],
    [INIT_VEL, 90], [INIT_VEL, 90], [INIT_VEL, 90], [INIT_VEL, 90]]
INIT_MASES = [m0, m1, m2, m3, m4, m5, m6, m7, m8]
########################################33
# Propiedaes para el algoritmo de flock
MAX_AGENTES = 30

# Giros maximo en grados
MAX_TURN_AVOID = 30#math.radians(.5)
MAX_TURN_COHE = 10#math.radians(5)
MAX_TURN_ALIGN = 10#math.radians(8)
# Separaciones minimas
MIN_SEP = 1.5*RADIO_AGENTES

INIT_AGENTS = [INIT_POS, INIT_VELO, INIT_MASES]

def main(init_agents=INIT_AGENTS, matriz=(70,70)):
    global cells, SIZE, CLEAR, RADIO_AGENTES, INIT_VEL, MIN_VEL, POS,VISION
    
    pygame.init
        
    filas = matriz[0]
    columnas = matriz[1]
    
    screen = pygame.display.set_mode((columnas*SIZE, filas*SIZE))
    
    background = pygame.Surface(screen.get_size())
    background = background.convert()
    background.fill( (255,255,255) )
    background_rect = background.get_rect()

    matrix = Matrix(filas, columnas, sizeCell=(SIZE,SIZE))
    grupo_cells = matrix.sprite_group

    clock = pygame.time.Clock()
    fps = 30
    
    # Creo paredes
    #grupo_pared = make_wall(matrix, (15,15),(30,15))
    #grupo_pared.add( make_wall(matrix, (15,15),(15,30)) )
    #grupo_pared.add( make_wall(matrix, (15,30),(22,30)) )
    #grupo_pared.add( make_wall(matrix, (30,30),(20,15)) )
    #grupo_pared.add( make_wall(matrix, (30,15),(30,30)) )
    #grupo_pared.add( make_wall(matrix, (30,30),(25,30)) )
    
    # Recupero la pantalla
    screen = pygame.display.get_surface()
    
    # Guardo area de la pantalla para hacer que salgan y entren
    area = screen.get_rect()
    
    # Creo a los agentes
    init_pos = init_agents[0]
    init_vel = init_agents[1]
    init_mases = init_agents[2]
    agentes = createAgents(init_pos, init_vel, init_mases)
    
    # add a fixed agent in the middle of the screen
    pos = [columnas*SIZE/2, filas*SIZE/2]
    masa = 10.0
    speed = [0,0]
    size_agent = (1,1)
    color = (0,0,0)
    #agente_aux = agente.Agente(pos=pos, masa=masa, \
    #              speed=speed, rangVision=0\
    #                ,size=size_agent,color=color)
    #agentes.append( agente_aux )
    grupo_agentes = pygame.sprite.Group(agentes)
    #grupo_agentes = pygame.sprite.Group()
    # Dibujo el fondo
    screen.blit(background, (0,0))
    
    # Dibujo a las celdas en mi screen
    grupo_cells.draw(screen)
    
    #grupo_agentes.update(grupo_pared)
    grupo_agentes.update()
    grupo_agentes.draw(screen)
    
    pygame.display.flip()
    
    tot_thresh = [0.02, 0.01];
    
    end = False
    count = 0
    first = True
    test_time = 0
    while (not end and test_time < 500):
        test_time += 1
        # Mantengo la actualizacion grafica a la velocidad deseada
        clock.tick(fps)
        count += 1
        
        # Llamo al manejador de eventos y obtengo posicion si se hizo click
        pos = event_handler()
        
        # Obtengo la celda donde se hizo click
        if len(pos) > 0 :
            createExplosion(matrix, pos)
            matrix.update()
        total = matrix.thresh
        
        if total > tot_thresh[0]:
            matrix.update()
        elif total > tot_thresh[1]:
            for x in range(0,matrix.rows+1):
                for y in range(0,matrix.cols+1):
                    matrix.setCellPressure(x, y, matrix.get_value(x,y)*0.85)
        else:
            # Actualizo la matriz con ceros
            matrix.setMatrixValues(0)
            
        for agent in grupo_agentes.sprites():
                
            ######################################
            #ACTUALIZACION DE AGENTES
            # 1. Hago el flocking
            flock(matrix, agent, agentes)

            ######################################
            #ACOPLAMIENTO CON CA
            if total > tot_thresh[1]:
                #f = fuerzaAcople(matrix, agent)
                f = fuerzaPresion(matrix, agent)
                
                # Actualizo la velocidad con la fuerza calculada
                agent.updateVel(f)
                
                # Si tiene mucha velocidad aplico una perdida de la misma
                if ( math.sqrt(agent.speed[0]**2+ agent.speed[1]**2) > MIN_VEL):
                    agent.speed[0] = agent.speed[0] * .8
                    agent.speed[1] = agent.speed[1] * .8               

        # the fixed agent
        #agentes[-1].speed = [0,0]
        
        # Dibujo el fondo
        screen.blit(background, (0,0))
        
        # Actualizo y dibujo las celdas
        grupo_cells.draw(screen)
        #grupo_pared.update()
        #grupo_pared.draw(screen)
        # Dibujo a los agentes
        #grupo_agentes.update(grupo_pared)
        grupo_agentes.update()
        grupo_agentes.draw(screen)
        
        # Muestro screen actualizado
        pygame.display.flip()
        
def fuerzaAcople(matrix, agent):
    # 2. Obtengo la celda en la posicion del agente
    aux_cell = matrix.get_cell(int(agent.rect.centery/SIZE) ,int(agent.rect.centerx/SIZE))
    
    # 3. Calculo Fuerza sobre el agente:
    box = agent.rect.left/SIZE, agent.rect.top/SIZE,agent.rect.right/SIZE, agent.rect.bottom/SIZE
    shw = matrix.getCellsBB(box)
    #print "Shadow: ",shw,shw.shape
    dP = util.cellGradient(shw)
    #   coordenadas x crecen hacia la derecha -> Este
    #   coordenadas y crecen hacia la abajo -> Sur
    dP = -dP
    #print "Pressure difference: ",dP
    
    if len(shw.shape) == 2:
        f = dP*shw.shape
    else:
        f = dP
    
    return f

def fuerzaPresion(matrix, agent):
    
    # 2. Obtengo la celda en la posicion del agente
    fila, columna = int(agent.rect.centery/SIZE) ,int(agent.rect.centerx/SIZE)
    cell = matrix.get_cell(fila ,columna)
    
    vecinos = matrix.vecinos_neumann(fila, columna)
    norte, sur, este, oeste = vecinos[0], vecinos[1], vecinos[2], vecinos[3]
    
    fx = este.get_pressure() - oeste.get_pressure()
    fy = norte.get_pressure() - sur.get_pressure()
    
    return fx, fy

def createExplosion(matrix, pos):
    for aux in pos :
        v = matrix.vecinos_neumann(int(aux[1]/float(SIZE)), int(aux[0]/float(SIZE)))
        for k in v:
            k.pressure = 0.045*matrix.rows*matrix.cols
            k.new_pressure = k.pressure
            
def createAgents(pos=None, vel=None, mases=None, color=(255,0,0)):
    '''Crea len(pos) agentes con las posiciones, velocidades y masas 
    pasadas como parametros (len(pos) = len(vel) = len(mases)).
    El array vel es de la forma vel = [[v_1,theta_1],[v_2,theta_2],...,[v_n,theta_n]]
    donde v_i es la magnitud del vector y theta_i es la amplitud en grados'''
    
    global MAX_AGENTES, VISION, RADIO_AGENTES
    size_agent = (RADIO_AGENTES,RADIO_AGENTES)
    agentes = []
    for i in range ( len(pos)):
        
        masa = mases[i]
        speed = [vel[i][0]*math.cos(math.radians(vel[i][1])),\
                vel[i][0]*math.sin(math.radians(vel[i][1])) ]
        
        agente_aux = agente.Agente(pos=pos[i], masa=masa, \
                speed=speed, rangVision=VISION\
                ,size=size_agent,color=color)
        agentes.append( agente_aux )
        
    return agentes


def flock(matrix, agent, agentes):
    
    global MAX_TURN_AVOID, MAX_TURN_COHE, MAX_TURN_ALIGN
    
    vision = pygame.sprite.GroupSingle(agent.vision)
    vision.draw(agent.screen)
    
    #vision = agente.Vision(agent.rect.center, 150)
    #pygame.sprite.spritecollide(sprite, group, dokill, collided = None): return Sprite_list
    
    neighbours = pygame.sprite.spritecollide(vision.sprites()[0], agentes, False, pygame.sprite.collide_circle)
    
    # a fixed agent comes in the last elment
    neighbours.append(agentes[-1])
    
    # Saco al agente de la lista de vecinos
    neighbours.remove(agent)

    #print agent.id,[n.id for n in neighbours]
        
    # Salgo si no hay vecinos
    if len(neighbours) <= 0:
        #move with noise
        agent.speed[0] += INIT_VEL*(2*random.randint(0,100)/100-1)
        agent.speed[1] += INIT_VEL*(2*random.randint(0,100)/100-1)
        return
        
    sep = 10000 
    x_velocities = []
    y_velocities = []
    x_positions = []
    y_positions = []
    mases = []
    distances = []
    for neighbour in neighbours:
        
        # Calculo obtengo la velocidad
        x_velocities.append(neighbour.speed[0])
        y_velocities.append(neighbour.speed[1])
        x_positions.append(neighbour.rect.centerx)
        y_positions.append(neighbour.rect.centery)
        mases.append( neighbour.masa )
        
        d = math.sqrt((neighbour.rect.centerx - agent.rect.centerx)**2 + (neighbour.rect.centery - agent.rect.centery)**2 )
        distances.append(d)
    
    # Obtengo la separacion minima y el vecino
    sep = min(distances)
    nearest = neighbours[distances.index(sep)]
    
    
    # Calculo el heading del agente
    theta = math.atan2(agent.speed[1], agent.speed[0])
    my_heading = math.degrees( theta )

    if ( my_heading < 0):
        my_heading += 360
    agent.heading = my_heading
    
    if (sep < MIN_SEP):
        # Aplico Separacion
        # Obtengo el signo de la vel
        turn_sep = util.separacion(agent, nearest)
        
        if turn_sep >= 0: 
            my_heading += min( turn_sep, MAX_TURN_AVOID)
        else:
            my_heading += max( turn_sep, -MAX_TURN_AVOID)
        
    else:
        # Aplico Alineamiento 
        turn_alig = util.alineacion(agent, x_velocities, y_velocities)
        
        if turn_alig >= 0:
            my_heading += min(turn_alig, MAX_TURN_ALIGN)
        else:
            my_heading += max(turn_alig, -MAX_TURN_ALIGN)
            
        # Aplico Cohesion
        turn_cohe, vel_media = util.cohesion(agent, x_positions, y_positions, x_velocities, y_velocities, mases)
        if turn_cohe >=0:
            my_heading += min( turn_cohe , MAX_TURN_COHE)
        else:
            my_heading += max( turn_cohe , -MAX_TURN_COHE)
            
        vel = vel_media
    # Direcciono la velocidad del agente con el angulo calculado
    
    vel = math.sqrt(agent.speed[0]**2 + agent.speed[1]**2)
    agent.speed[0] += 5*(math.cos(math.radians(my_heading)) - math.cos(math.radians(agent.heading))) 
    agent.speed[1] += 5*(math.sin(math.radians(my_heading)) - math.sin(math.radians(agent.heading))) 
    #vision.kill
    return


def event_handler():
    global DOWN, CLEAR
    LEFT = 1
    RIGHT = 3
    ret = []
    
    for event in pygame.event.get():
        
        if event.type == pygame.QUIT:
            sys.exit(0)
        elif event.type == pygame.MOUSEBUTTONDOWN and event.button == LEFT:
            DOWN = True
            ret.append(event.pos)
        elif event.type == pygame.MOUSEBUTTONUP:
            DOWN = False
        #        event = pygame.event.poll()
        elif event.type == pygame.MOUSEMOTION and DOWN:
            ret.append(event.pos)
        elif event.type == pygame.MOUSEBUTTONDOWN and event.button == RIGHT:
            CLEAR = True
    
    return ret
 
def make_wall(matrix, p_from, p_to, color=(255,0,0)):
    pared_cells = pygame.sprite.Group()
    turn = False
    if (p_from[0] == p_to[0]):
        '''Hago linea paralela a eje Y desde Y0 hasta Y1'''
        #make_wall_y (p_from[1], p_to[1])
        x0, y0 = p_from[1], p_from[0]
        x1, y1 = p_to[1], p_to[0]
        turn = True
    else:
        x0, y0 = p_from[0], p_from[1]
        x1, y1 = p_to[0], p_to[1]

    #print "Creating wall from (%s, %s) to (%s, %s)"%(x_ini,y_ini, x_fin, y_fin)
    m, b = get_fun_param ((x0,y0), (x1,y1))
  
    for x in range(x0, x1):
        
        y = int(m*x + b)
        if turn:
            #print " celda : (%s, %s)"%(y,x)
            cell = matrix.get_cell(y,x)
        else:
            #print " celda : (%s, %s)"%(x,y)
            cell = matrix.get_cell(x,y)
        pared_cells.add( cell )
        cell.allow_flow = False
        cell.image.fill( color )
        
    return pared_cells

def get_fun_param(p0, p1):
    ''' f(x) = m*x + b. Esta funcion devuelve los parametros m, b para p0 y p1'''
    dist_x = (p1[0] - p0[0] )
    
    if not dist_x == 0:
        m = (p1[1] - p0[1]) / dist_x
    else:
        return None, None
    
    b = p1[1] - m*p1[0]
    
    return m, b
    
if __name__ == "__main__":
    main()
